package com.chenfan.gateway.common.utils;

import org.slf4j.MDC;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.web.server.ServerWebExchange;

import java.util.Objects;
import java.util.UUID;

/**
 * traceid工具类
 */
public class TraceUtil {
    public static final String TRACE_ID = "traceId";
    public static String getUUID() {
        return UUID.randomUUID().toString().replace("-", "");
    }

    public static void logMDC(ServerWebExchange exchange, String key, String value) {
        if (Objects.isNull(key) || Objects.isNull(value)){
            return;
        }
        //对请求对象request进行增强
        ServerHttpRequest request = exchange.getRequest();
        ServerHttpResponse response = exchange.getResponse();

        request.mutate().headers(httpHeaders -> {
            //httpHeaders 封装了所有的请求头
            String oldValue = request.getHeaders().getFirst(key);
            String newValue = Objects.isNull(oldValue) ? value : oldValue;
            MDC.put(key, newValue);
            httpHeaders.set(key, newValue);

        }).build();

        ServerWebExchange.Builder builder = exchange.mutate().request(request);
        // 增加 traceId
        if(TRACE_ID.equals(key)){
            String traceId = response.getHeaders().getFirst(TRACE_ID);
            if(Objects.isNull(traceId)){
                response.getHeaders().set(key,value);
                builder.response(response);
            }
        }

    }
}
